# -*- coding: utf-8 -*-

'''
基于TF-IDF提取关键词
'''

import jieba.posseg as pseg
from sklearn.feature_extraction.text import TfidfVectorizer
from collections import Counter
import re

class KeywordExtractor:
    def __init__(self, stopwords_path=None, custom_stopwords=None, min_freq=2, max_freq=50, allowed_pos=None):
        """
        初始化关键词提取器。
        :param stopwords_path: 停用词文件路径
        :param custom_stopwords: 自定义干扰词列表
        :param min_freq: 词频过滤的最小频率
        :param max_freq: 词频过滤的最大频率
        :param allowed_pos: 允许的词性列表，如 {'n', 'ns'} 表示名词和地名
        """
        self.stopwords = set()
        if stopwords_path:
            self.load_stopwords(stopwords_path)
        if custom_stopwords:
            self.stopwords.update(custom_stopwords)
        self.min_freq = min_freq
        self.max_freq = max_freq
        self.allowed_pos = allowed_pos

    def load_stopwords(self, filepath):
        """加载停用词表"""
        with open(filepath, 'r', encoding='utf-8') as f:
            self.stopwords = set(line.strip() for line in f)

    def clean_text(self, text):
        """清理文本，去除非中文字符"""
        text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
        return text

    def tokenize(self, text):
        """分词、过滤停用词，并按词性筛选"""
        words = pseg.cut(text)
        filtered_words = [
            word for word, flag in words
            if word not in self.stopwords and (self.allowed_pos is None or flag in self.allowed_pos)
        ]
        return filtered_words

    def filter_by_frequency(self, words):
        """根据词频过滤"""
        word_counts = Counter(words)
        return [word for word in words if self.min_freq <= word_counts[word] <= self.max_freq]

    def extract_keywords_tfidf(self, documents, top_k=20):
        """
        使用TF-IDF提取关键词
        :param documents: 文本列表
        :param top_k: 每篇文档返回的关键词数量
        :return: 每个文档的关键词列表
        """
        cleaned_docs = [" ".join(self.tokenize(self.clean_text(doc))) for doc in documents]
        tfidf_vectorizer = TfidfVectorizer()
        tfidf_matrix = tfidf_vectorizer.fit_transform(cleaned_docs)
        feature_names = tfidf_vectorizer.get_feature_names_out()

        keywords = []
        for doc_idx in range(tfidf_matrix.shape[0]):
            tfidf_scores = tfidf_matrix[doc_idx].toarray()[0]
            word_score_pairs = list(zip(feature_names, tfidf_scores))
            sorted_words = sorted(word_score_pairs, key=lambda x: x[1], reverse=True)
            doc_keywords = [word for word, score in sorted_words[:top_k]]
            keywords.append(doc_keywords)

        return keywords

    def extract_keywords(self, text, top_k=20):
        """
        综合过滤后的关键词提取
        :param text: 文本字符串
        :param top_k: 返回的关键词数量
        :return: 关键词列表
        """
        # 清理和分词
        cleaned_text = self.clean_text(text)
        words = self.tokenize(cleaned_text)
        
        # 词频过滤
        words = self.filter_by_frequency(words)
        
        # 提取最重要的关键词
        tfidf_keywords = self.extract_keywords_tfidf([text], top_k=top_k)
        return tfidf_keywords[0]
    
if __name__ == '__main__':
    path = "/home/ubuntu/code/git/subject-word-extraction/scripts/clean_data/000001_2023_平安银行_2023年年度报告_2024-03-15.txt"
    with open(path,"r") as f:
        content = f.read()

    # 初始化关键词提取器
    extractor = KeywordExtractor(
        stopwords_path="/home/ubuntu/code/git/subject-word-extraction/scripts/stop_word/stopwords/cn_stopwords.txt",
        custom_stopwords={"股份有限公司", "年度报告", "财务数据", '资产', '集团', '银行', '风险', '有限公司', '股份', '年度', '价值', '人民币', '金融',
                        '业务', '负债', '财务报表', '金融资产', '金额', '贷款', '计量', '交易', '情况', '单位', '收益', '客户', '损失', '项目', '余额',
                        '中国', '信用', '企业', '董事', '股东', '损益', '报告', '市场', '公司', '资金', '垫款', '资本', '利率', '权益', '成本', '现金', 
                        '金融工具', '款项', '债权', '工具', '利息', '日本', '发放贷款', '计划', '存款', '能力', '债券', '产品', '流动性', '所得税', '信用风险', 
                        '信息', '现金流量', '董事会', '风险管理', '账面', '金融机构', '机构', '证券', '经济', '总额', '管理层', '产生', '事项', '优先股', 
                        '货币', '方面', '政策', '主体', '数据', '委员会', '债务', '手续费', '先生', '净利润', '模型', '营业', '分部', '费用', '比例', '行业', 
                        '分类', '普通股', '国家', '平台', '个人', '会议', '体系', '实际', '深圳', '部分', '子公司', '净额', '现任', '结构', '负债表', '质量', 
                        '摊余', '结构化', '股本', '深圳市', '监事', '实体', '本金', '薪酬', '财务', '战略', '阶段', '关联', '套期', '会计师', '关键', '核心', 
                        '股权', '利润', '外币', '数量', '方式', '核销', '金融服务', '办法', '条件', '净收入', '模式', '管理人员', '范围', '计提', '重点', '基础', 
                        '客群', '规定', '领域', '全面', '政府', '行长', '利息收入', '高质量', '使用权','期限', '股东大会', '公积', '助力', '投资者', '指标', '标准', '财务报告', '责任', '佣金', '差异', '事务所', '力度', '商誉', '地区', '对公', '贵金属', '专业', '充足率', '利润分配', '基本', '技术', '职工', 
                        '副行长', '场景', '方法', '监督管理', '评价', '义务', '借款人', '关系', '国际', '审计报告', '水平', '社会', '资产负债', '中心', '信托', '机制', 
                        '优势', '信用卡', '方案', '股东权益', '股息', '规模', '投资性', '母公司', '基金', '平安保险', '总行', '报酬', '收益率', '整体', '系统', '账户', '因素', '定价', '总经理', '报告期', '流程', '环境', '股利', '会计准则', '国别', '委托', '总局', '汇率', '参数', '增加额', '盈余', '账簿', '净资产', 
                        '员工', '法人', '监事会', '借款', '名称', '工程', '差额', '消费者', '目标', '问题', '首席', '人寿保险', '全行', '制度', '情景', '票据', '策略', '纳税', 
                        '事件', '产业', '体验', '关联方', '净值', '外汇', '管理体系', '路号', '代理', '估值', '使用寿命', '存续期', '措施', '概率', '法定', '跨境', '副总经理', '文化', '渠道', '目的', '福利', '类型', '议案', '分配利润', '委员', '宏观经济', '意见', '投资收益', '教授', '架构', '法律', '经常性', '财富', '金融业务', 
                        '限额', '不良贷款', '人员', '价格', '前瞻性', '历任', '报价', '时本', '状况', '理财产品', '董事长', '试验区', '部门', '错报', '队伍', '上海', '住房', '利润表', '原因', '效率', '无法', '程序', '资源', '优质', '大厦', '布局', '成果', '担保物', '方方', '特征', '现金流', '票面', '衍生品', '覆盖率', '不良率', '主席', '利益',
                        '声誉', '大道', '导向', '寿险', '年本', '损失率', '时间', '权利', '经济学', '财务指标', '业绩', '任期', '保险', '党委书记', '公积金', '内容', '功能', '协议', '原则', '含税', '四本', '国债', '律师', '有效性', '权重', '深度', '职务', '预案', '专项', '东区', '东路', '借贷', '决策', '压力', '女士', '性质', '政策性', '条款', '程度', '管理部', '职责', '重要性', '银行业', '上市', '中央', '办公', '区域'
                        '品牌', '地址', '基准', '姓名', '广东省', '建筑物', '总体', '总部', '息差', '法律法规', '特色', '精准', '股票', '规划', '计息', '额度', '上线', '专门', '北区', '协同', '南区', '发生额', '奖金', '存量', '房屋', '有志', '核准', '次数', '比率', '状态', '理念', '电话', '私行', '西区', '过程', 
                        '书记', '代表', '价差', '保函', '信用证', '公司债券', '公告', '办公室', '厦门', '口径', '基点', '学院', '待摊费用', '所得额', '敏感性', '案件', '源生', '租赁期', '管理策略', '网点', '落地', '解决方案', '购买方', '金融中心', '面值', '香港', '专家', '业务收入', '人才', '保险产品', '信用等级', '全额'
                        '高风险', '准备金', '动态', '原值', '名义', '名单', '团队', '基数', '天津', '头寸', '广东', '广场', '形势', '总数', '成员', '执行官', '报表', '指导', '新能源', '武汉', '武汉市', '法定代表', '注册资本', '福田区', '税率', '税费', '管理工作', '精神', '经营性', '缺口', '股份制', '行使', '账龄'
                        '全额', '院长', '青岛', '领先', '专业化', '中华人民共和国', '代理人', '份额', '传统', '保持良好', '办理', '博士', '口袋', '合同条款', '地方', '城市', '大学', '客户服务', '导师', '建议', '强国', '形式', '总计', '总资产', '承兑汇票', '控制权', '权力', '杠杆', '案例', '民营企业', '深南', '用户'
                        '高风险', '全额', '账龄', '电子', '盈利', '福建', '私人', '租赁费', '等本', '管理制度', '管理机制', '类别', '纠纷', '薪金', '规则', '贷后', '资料', '金融债', '银行家', '集中度', '中路', '产业链', '准则', '分支机构', '分派', '利润总额', '前提', '北京', '北京市', '区分', '南路', '可能性', '合法权益'
                        '附注', '分行', '减值', '准备', '相关', '计入', '注明', '适用', '进行', '确认', '控制', '发行', '合并', '提升', '包括', '衍生', '会计', '披露', '发生', '存放', '执行', '提供', '收入', '增加', '支付', '支持', '独立', '减少', '增长', '是否', '审计', '要求', '承诺', '租赁', '确定', '治理', '控股', '违约', '存在', '加强', '处置', '支出', '强化', '使用', '持有', '吸收', '递延', '采用'
                        '实现', '导致', '信贷', '计算', '转入', '完善', '累计', '可能', '下降', '创新', '开展', '拆出', '初始', '审议', '推进', '估计', '讨论', '理财', '判断', '回购', '建立', '打造', '敞口', '买入', '满足', '终止', '返售', '列示', '获得', '数字化', '取得', '应计', '考虑', '赋能', '坚持', '增减', '加大', '担保', '批准', '结算', '履行', '加权', '实施', '应收', '推动',
                        '附注', '采用', '实现', '卖出', '承担', '拆入', '摊销', '纳入', '折现', '数字', '聚焦', '需求', '变更', '形成', '折旧', '注册', '后续', '收回', '符合', '说明', '分红', '参与', '设定', '具有', '收到', '组合', '不能', '反映', '扣除', '持股', '操作', '转换', '上升', '修订', '相应', '负责', '投诉', '暴露', '能够', '保持', '包含', '发放', '应用', '托管',
                        '精细化', '区域', '品牌', '金融债券', '金融市场', '高风险', '用户', '合法权益', '合理性', '商业', '商品', '土地', '大额', '宏观', '定义', '对价', '工会', '币种', '年限', '底座', '总监', '指数', '指本', '文件', '有序', '有权', '服务平台', '机会', '津贴', '溢价', '现值', '生态圈', '监管部门', '硕士', '税金', '行通', '补贴', '规范', '账款', '资产重组', '迹象'},
        min_freq=5,
        max_freq=50,
        allowed_pos={'n', 'ns', 'v'}
    )

    # 提取关键词
    keywords = extractor.extract_keywords(' '.join(content), top_k=None)
    print("提取的关键词：", keywords)